import cv2
import numpy as np
import os

def generate_handwritten_digit(digit, filename=None):
    """
    生成28x28的手写数字图片

    参数:
    digit: 要生成的数字 (0-9)
    filename: 保存的文件名，如果不提供则显示图片但不保存
    """

    # 创建28x28的黑色背景图像
    img = np.zeros((28, 28), dtype=np.uint8)

    # 设置字体和参数
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 1
    thickness = 2
    color = 255  # 白色

    # 获取数字的大小
    text_size = cv2.getTextSize(str(digit), font, font_scale, thickness)[0]

    # 计算文字位置使其居中
    text_x = (28 - text_size[0]) // 2
    text_y = (28 + text_size[1]) // 2

    # 在图像上绘制数字
    cv2.putText(img, str(digit), (text_x, text_y), font, font_scale, color, thickness)

    # 添加一些手写效果
    # 添加轻微的扭曲
    noise = np.random.randint(0, 30, (28, 28), dtype=np.uint8)
    img = cv2.add(img, noise)

    # 应用高斯模糊使边缘更柔和
    img = cv2.GaussianBlur(img, (3, 3), 0)

    # 显示图像
    cv2.imshow(f'Digit {digit}', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 如果提供了文件名，则保存图像
    if filename:
        # 确保目录存在
        directory = "../data/img"
        if not os.path.exists(directory):
            os.makedirs(directory)

        # 构建完整路径
        full_path = os.path.join(directory, filename)
        cv2.imwrite(full_path, img)
        print(f"图片已保存为: {full_path}")

    return img

def generate_multiple_digits():
    """
    生成0-9所有数字的图片并保存在../data/img/目录下
    """
    directory = "../data/img"
    if not os.path.exists(directory):
        os.makedirs(directory)

    for i in range(10):
        filename = f"digit_{i}.png"
        full_path = os.path.join(directory, filename)
        generate_handwritten_digit(i, full_path)

if __name__ == "__main__":
    # 输入要生成的数字
    try:
        digit = int(input("请输入要生成的数字 (0-9): "))
        if 0 <= digit <= 9:
            filename = f"digit_{digit}.png"
            img = generate_handwritten_digit(digit, filename)
            print(f"已生成数字 {digit} 的图片")
        else:
            print("请输入0-9之间的数字")
    except ValueError:
        print("请输入有效的数字")
